
import chisel3._
import chiseltest._
import org.scalatest.{Matchers, FlatSpec}
import scala.util.Random

import comlib._
import comlib.util.{reverse}
import comlib.arbiter._


class EArbiterSpec extends FlatSpec with ChiselScalatestTester with Matchers {

  it should "support partial assign for input signals" in {
    test( new Module{
      val io = IO(new Bundle{
        val rmi = Input(UInt(4.W))
        val cmi = Input(UInt(4.W))
        val gr = Output(Vec(4, Vec(4, Bool())))
      })

      val arb = Module( new EArbiter(4, 4) )
      arb.io.rmi(0) := io.rmi.asBools 
      arb.io.cmi(0) := io.cmi.asBools
      arb.io.req(0) := VecInit( Seq.fill(4)( VecInit( Seq.fill(4)(true.B) ) ) )
      io.gr := arb.io.gr0

    } ){ dut => 

      dut.io.rmi.poke("b1010".U)
      dut.io.cmi.poke("b0101".U)
      for( r <- 0 until 4; c <- 0 until 4 ){
        if( r==0&&c==1 || r==2&&c==3 ) dut.io.gr(r)(c).expect( true.B )
        else dut.io.gr(r)(c).expect( false.B )
      }

    }
  }


  it should "support rmi mask" in {
    test( new EArbiter(4, 4) ){ dut => 
      dut.io.cmi(0).foreach { case cm => cm.poke(false.B) }
      dut.io.req(0).foreach { case reqr => reqr.foreach { case reqc => reqc.poke(true.B) } }

      for( it <- 0 until 4 ){
        dut.io.rmi(0).zipWithIndex.foreach { case (sig, r) => sig.poke( (r==it).B ) }

        for( r <- 0 until 4; c <- 0 until 4 ){
          val hasGnt = r<it&&c==r || r>it&&c==r-1
          dut.io.gr0(r)(c).expect( hasGnt.B, f"it=${it} r=${r} c=${c}" )
          dut.io.gc0(c)(r).expect( hasGnt.B, f"it=${it} r=${r} c=${c}" )
        }

        dut.io.grv0.zipWithIndex.foreach { case (sig, r) => sig.expect((r!=it).B, f"it=${it} r=${r}") }
        dut.io.rmo0.zipWithIndex.foreach { case (sig, r) => sig.expect(true.B) }
        dut.io.gcv0.zipWithIndex.foreach { case (sig, c) => sig.expect((c!=3).B, f"it=${it} c=${c}") }
        dut.io.cmo0.zipWithIndex.foreach { case (sig, c) => sig.expect((c!=3).B, f"it=${it} c=${c}") }

      }
    }
  }

  it should "support cmi mask" in {
    test( new EArbiter(4, 4) ){ dut => 
      dut.io.rmi(0).foreach { case cm => cm.poke(false.B) }
      dut.io.req(0).foreach { case reqr => reqr.foreach { case reqc => reqc.poke(true.B) } }

      for( it <- 0 until 4 ){
        dut.io.cmi(0).zipWithIndex.foreach { case (sig, r) => sig.poke( (r==it).B ) }

        for( r <- 0 until 4; c <- 0 until 4 ){
          val hasGnt = c<it&&r==c || c>it&&r==c-1
          dut.io.gc0(c)(r).expect( hasGnt.B, f"it=${it} r=${r} c=${c}" )
          dut.io.gr0(r)(c).expect( hasGnt.B, f"it=${it} r=${r} c=${c}" )
        }

        dut.io.gcv0.zipWithIndex.foreach { case (sig, r) => sig.expect((r!=it).B, f"it=${it} r=${r}") }
        dut.io.cmo0.zipWithIndex.foreach { case (sig, r) => sig.expect(true.B) }
        dut.io.grv0.zipWithIndex.foreach { case (sig, c) => sig.expect((c!=3).B, f"it=${it} c=${c}") }
        dut.io.rmo0.zipWithIndex.foreach { case (sig, c) => sig.expect((c!=3).B, f"it=${it} c=${c}") }

      }
    }
  }

  it should "support grv&gcv output" in {
    test( new Module{
      val io = IO(new Bundle{
        val grv = Output(Vec(4, Bool()))
        val gcv = Output(Vec(4, Bool()))
      })

      val arb = Module( new EArbiter(4, 4) )
      arb.io.rmi(0) := VecInit( false.B, true.B, false.B, true.B )
      arb.io.cmi(0) := VecInit( true.B, false.B, true.B, false.B )
      // arb.io.req(0) := VecInit( Seq.fill(4)( VecInit( Seq.fill(4)(true.B) ) ) )
      arb.io.req := arb.req(true.B)
      io.grv := arb.io.grv0 
      io.gcv := arb.io.gcv0

    } ){ dut => 

      dut.io.grv.zipWithIndex.foreach { case (g,i) => if( i==0 || i==2 ) g.expect(true.B) else g.expect(false.B) }
      dut.io.gcv.zipWithIndex.foreach { case (g,i) => if( i==1 || i==3 ) g.expect(true.B) else g.expect(false.B) }

    }
  }

  
  it should "support SP/SPReverse scheduler for SelectFind1 and Select1Parrel Action" in {

    class SpModule(action:SelectAction, isRev:Boolean) extends Module{
      val io = IO(new Bundle{
        val req = Input(UInt(4.W))
        val gnt = Output(UInt())
        val gv = Output(Bool())
      })
      val arb = if(isRev) EArbiter.SPReverse(io.req, action) else EArbiter.SP(io.req, action)
      io.gnt := arb.io.gnt.asUInt
      io.gv := arb.io.gv
    }

    val parameters = Seq( new SelectFind1, new Select1Parrel )

    for( action <- parameters; rev <- 0 until 1 ){
      test( new SpModule(action, rev==1) ){ dut =>

        val result = Seq(
          "b1111".U -> ("b0001".U, true.B),
          "b1110".U -> ("b0010".U, true.B),
          "b1100".U -> ("b0100".U, true.B),
          "b1000".U -> ("b1000".U, true.B),
          "b0000".U -> ("b0000".U, false.B),
          "b1010".U -> ("b0010".U, true.B),
        )

        for( (r, (g, gv)) <- result ){
          val req = if(rev==1) reverse(r) else r 
          val gnt = if(rev==1) reverse(g) else g
          dut.io.req.poke(req)
          dut.io.gnt.expect(gnt, f"req=${req} gnt=${gnt} action=${action}")
          dut.io.gv.expect(gv, f"req=${req} gnt=${gnt} action=${action}")
        }

      }
    }

  }

  it should "support 1D RR scheduler" in {
    test( new Module{
      val io = IO(new Bundle{
        val req = Input(UInt(4.W))
        val ptr = Input(UInt(4.W))
        val gnt = Output(UInt(4.W))
        val gv = Output(Bool())
      })
      val arb = EArbiter.RoundRobin( VecInit(io.req.asBools), VecInit(io.ptr.asBools) )
      io.gnt := arb.io.gnt.asUInt
      io.gv := arb.io.gv
    } ){ dut =>

      val result = Seq(
        ("b1111".U, "b0000".U) -> ("b0001".U, true),
        ("b1111".U, "b0001".U) -> ("b0010".U, true),
        ("b1111".U, "b0011".U) -> ("b0100".U, true),
        ("b1111".U, "b0111".U) -> ("b1000".U, true),
        ("b1100".U, "b0000".U) -> ("b0100".U, true),
        ("b1110".U, "b0011".U) -> ("b0100".U, true),
        ("b0011".U, "b0011".U) -> ("b0001".U, true),
        ("b0110".U, "b0111".U) -> ("b0010".U, true),
      )

      for( ((req, ptr), (gnt, gv)) <- result ){
        dut.io.req.poke(req)
        dut.io.ptr.poke(ptr)
        dut.io.gnt.expect(gnt)
        dut.io.gv.expect(gv.B)
      }

    }
  }

  it should "support 1D Unitl0 scheduler" in {
    test( new Module{
      val io = IO(new Bundle{
        val req = Input(UInt(4.W))
        val ptr = Input(UInt(4.W))
        val gnt = Output(UInt(4.W))
        val gv = Output(Bool())
      })
      val arb = EArbiter.Until0( VecInit(io.req.asBools), VecInit(io.ptr.asBools) )
      io.gnt := arb.io.gnt.asUInt
      io.gv := arb.io.gv
    } ){ dut =>

      val result = Seq(
        ("b0000".U, "b0000".U) -> ("b0000".U, false),
        ("b0001".U, "b0000".U) -> ("b0001".U, true),
        ("b0011".U, "b0000".U) -> ("b0011".U, true),
        ("b0111".U, "b0000".U) -> ("b0111".U, true),
        ("b1111".U, "b0000".U) -> ("b1111".U, true),
        ("b1111".U, "b0001".U) -> ("b1111".U, true),
      )

      for( ((req, ptr), (gnt, gv)) <- result ){
        dut.io.req.poke(req)
        dut.io.ptr.poke(ptr)
        dut.io.gnt.expect(gnt, f"req=${req} ptr=${ptr} gnt=${gnt} gv=${gv}")
        // dut.io.gv.expect(gv.B)
      }

    }
  }

  it should "support RRDPA arbiter" in {

    class MyModule extends Module{
      val io = IO(new Bundle{
        val req = Input(Vec(4, UInt(4.W)))
        val ptr = Input(UInt(4.W))
        val gnt = Output(Vec(4, UInt(4.W)))
      })
      val arb = EArbiter.RRDPA( io.req, io.ptr, 0.U, 0.U )
      io.gnt := VecInit(arb.io.gr0.map{ case g => g.asUInt })
    }

    val result = Seq(
      ("b1111".U, "b0001".U) -> Seq("b0001".U, "b0010".U, "b0100".U, "b1000".U),
      ("b1111".U, "b0010".U) -> Seq("b1000".U, "b0001".U, "b0010".U, "b0100".U),
      ("b1111".U, "b0100".U) -> Seq("b0100".U, "b1000".U, "b0001".U, "b0010".U),
      ("b1111".U, "b1000".U) -> Seq("b0010".U, "b0100".U, "b1000".U, "b0001".U),
      ("b0011".U, "b0001".U) -> Seq("b0001".U, "b0010".U, "b0000".U, "b0000".U),
      ("b0011".U, "b0010".U) -> Seq("b0000".U, "b0001".U, "b0010".U, "b0000".U),
      ("b0011".U, "b0100".U) -> Seq("b0000".U, "b0000".U, "b0001".U, "b0010".U),
      ("b0011".U, "b1000".U) -> Seq("b0010".U, "b0000".U, "b0000".U, "b0001".U)
    )

    val actions = Seq(new SelectFind1, new Select1Parrel)
    
    for( act <- actions ){
      test( (new MyModule).setAction(act) ){ dut => 

        for( ((req,ptr), gnt) <- result ){
          dut.io.req.foreach{ case r => r.poke(req) }
          dut.io.ptr.poke(ptr)
          (0 until 4).foreach{ case r => dut.io.gnt(r).expect(gnt(r)) }
        }
      }
    }

  }

  it should "support RoundRobinReverse Scheduler" in {

    class MyModule extends Module{
      val io = IO(new Bundle{
        val req = Input(UInt(4.W))
        val ptr = Input(UInt(4.W))
        val gnt = Output(UInt(4.W))
      })

      val arb = EArbiter.RoundRobinReverse( VecInit(io.req.asBools), VecInit(io.ptr.asBools) )
      io.gnt := arb.io.gnt.asUInt
    }

    test(new MyModule){ dut => 
      val result = Seq(
        ("b0000".U, "b0000".U) -> "b0000".U,    // 无请求，输出0
        ("b0000".U, "b0001".U) -> "b0000".U,    // 无请求，指针任意，输出0
        ("b0000".U, "b0011".U) -> "b0000".U,    // 无请求，指针任意，输出0
        ("b1111".U, "b0000".U) -> "b0001".U,    // 全请求，输出指针位置
        ("b1111".U, "b0001".U) -> "b0010".U,    // 全请求，输出指针位置
        ("b0111".U, "b0111".U) -> "b0100".U,    // 指针指到最后，最后无请求。输出前一个
        ("b0011".U, "b0111".U) -> "b0010".U,    // 指针指到最后，最后无请求。输出前一个
        ("b0001".U, "b0111".U) -> "b0001".U,    // 指针指到最后，最后无请求。输出前一个
        ("b1000".U, "b0011".U) -> "b1000".U,    // 指针指到最后，最后无请求。输出前一个
        ("b0110".U, "b0000".U) -> "b0100".U,    // 指针指到最后，最后无请求。输出前一个
        ("b1110".U, "b0000".U) -> "b1000".U,    // 指针指到最后，最后无请求。输出前一个
      )

      for( ((req, ptr), gnt) <- result ){
        dut.io.req.poke(req)
        dut.io.ptr.poke(ptr)
        dut.io.gnt.expect(gnt, f"req=${req} ptr=${ptr} gnt=${gnt}")
        dut.clock.step()
      }
    }

  }

}
